{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e3eaf6be",
   "metadata": {},
   "source": [
    "## Enhanced Sidekick with Adaptive Clarification\n",
    "\n",
    "### Personal Co-worker with Planner Node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e510bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "import asyncio\n",
    "import requests\n",
    "import nest_asyncio\n",
    "from typing import Annotated\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from dotenv import load_dotenv\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "from typing import List, Any, Optional, Dict\n",
    "from pydantic import BaseModel, Field\n",
    "from datetime import datetime\n",
    "from playwright.async_api import async_playwright\n",
    "from langchain_community.agent_toolkits import PlayWrightBrowserToolkit\n",
    "from langchain.agents import Tool\n",
    "from langchain_community.agent_toolkits import FileManagementToolkit\n",
    "from langchain_community.tools.wikipedia.tool import WikipediaQueryRun\n",
    "from langchain_experimental.tools import PythonREPLTool\n",
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "from langchain_community.utilities.wikipedia import WikipediaAPIWrapper\n",
    "import gradio as gr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dcbe712",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9184455c",
   "metadata": {},
   "source": [
    "### Define State / Structured Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ff5870d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    messages: Annotated[List[Any], add_messages]\n",
    "    success_criteria: str\n",
    "    feedback_on_work: Optional[str]\n",
    "    success_criteria_met: bool\n",
    "    user_input_needed: bool\n",
    "    clarifying_questions_asked: int\n",
    "    planning_complete: bool\n",
    "\n",
    "\n",
    "class EvaluatorOutput(BaseModel):\n",
    "    feedback: str = Field(description=\"Feedback on the assistant's response\")\n",
    "    success_criteria_met: bool = Field(description=\"Whether the success criteria have been met\")\n",
    "    user_input_needed: bool = Field(\n",
    "        description=\"True if more input is needed from the user, or clarifications, or the assistant is stuck\"\n",
    "    )\n",
    "\n",
    "\n",
    "class PlannerOutput(BaseModel):\n",
    "    clarification_question: Optional[str] = Field(\n",
    "        description=\"A single clarifying question to ask the user, or None if everything is clear\"\n",
    "    )\n",
    "    ready_to_proceed: bool = Field(\n",
    "        description=\"True if the request and success criteria are clear enough to proceed to the worker\"\n",
    "    )\n",
    "    reasoning: str = Field(\n",
    "        description=\"Brief explanation of why clarification is needed or why ready to proceed\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "279e3fdd",
   "metadata": {},
   "source": [
    "### Sidekick Class with Planner, Worker, and Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "536fc826",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sidekick:\n",
    "    def __init__(self):\n",
    "        self.worker_llm_with_tools = None\n",
    "        self.evaluator_llm_with_output = None\n",
    "        self.planner_llm_with_output = None\n",
    "        self.tools = None\n",
    "        self.graph = None\n",
    "        self.sidekick_id = str(uuid.uuid4())\n",
    "        self.memory = MemorySaver()\n",
    "        self.browser = None\n",
    "        self.playwright = None\n",
    "\n",
    "    async def setup(self):\n",
    "        self.tools = await self.get_all_tools()\n",
    "        worker_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "        self.worker_llm_with_tools = worker_llm.bind_tools(self.tools)\n",
    "        evaluator_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "        self.evaluator_llm_with_output = evaluator_llm.with_structured_output(EvaluatorOutput)\n",
    "        planner_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "        self.planner_llm_with_output = planner_llm.with_structured_output(PlannerOutput)\n",
    "        await self.build_graph()\n",
    "\n",
    "    async def get_all_tools(self):\n",
    "        self.playwright = await async_playwright().start()\n",
    "        self.browser = await self.playwright.chromium.launch(headless=False)\n",
    "        toolkit = PlayWrightBrowserToolkit.from_browser(async_browser=self.browser)\n",
    "        browser_tools = toolkit.get_tools()\n",
    "\n",
    "        pushover_token = os.getenv(\"PUSHOVER_TOKEN\")\n",
    "        pushover_user = os.getenv(\"PUSHOVER_USER\")\n",
    "        pushover_url = \"https://api.pushover.net/1/messages.json\"\n",
    "\n",
    "        def push(text: str):\n",
    "            print(f\"[PUSH] Attempting to send notification: {text[:50]}...\")\n",
    "            if not pushover_token or not pushover_user:\n",
    "                error_msg = \"Notification not configured - missing PUSHOVER_TOKEN or PUSHOVER_USER in environment\"\n",
    "                print(f\"[PUSH ERROR] {error_msg}\")\n",
    "                return error_msg\n",
    "            print(\"[PUSH] Sending to Pushover API...\")\n",
    "            try:\n",
    "                response = requests.post(pushover_url, data={\"token\": pushover_token, \"user\": pushover_user, \"message\": text})\n",
    "                print(f\"[PUSH] Response status: {response.status_code}\")\n",
    "                if response.status_code == 200:\n",
    "                    success_msg = \"Notification sent successfully\"\n",
    "                    print(f\"[PUSH SUCCESS] {success_msg}\")\n",
    "                    return success_msg\n",
    "                error_msg = f\"Notification failed: {response.status_code} - {response.text}\"\n",
    "                print(f\"[PUSH ERROR] {error_msg}\")\n",
    "                return error_msg\n",
    "            except Exception as e:\n",
    "                error_msg = f\"Notification exception: {str(e)}\"\n",
    "                print(f\"[PUSH EXCEPTION] {error_msg}\")\n",
    "                return error_msg\n",
    "\n",
    "        push_tool = Tool(\n",
    "            name=\"send_push_notification\",\n",
    "            func=push,\n",
    "            description=\"Use this tool when you want to send a push notification\"\n",
    "        )\n",
    "\n",
    "        file_toolkit = FileManagementToolkit(root_dir=\"sandbox\")\n",
    "        file_tools = file_toolkit.get_tools()\n",
    "\n",
    "        serper = GoogleSerperAPIWrapper()\n",
    "        search_tool = Tool(\n",
    "            name=\"search\",\n",
    "            func=serper.run,\n",
    "            description=\"Use this tool when you want to get the results of an online web search\"\n",
    "        )\n",
    "\n",
    "        wikipedia = WikipediaAPIWrapper()\n",
    "        wiki_tool = WikipediaQueryRun(api_wrapper=wikipedia)\n",
    "\n",
    "        python_repl = PythonREPLTool()\n",
    "\n",
    "        mailgun_api_key = os.getenv(\"MAILGUN_API_KEY\")\n",
    "        mailgun_domain = os.getenv(\"MAILGUN_DOMAIN\")\n",
    "        mailgun_from = os.getenv(\"MAILGUN_FROM_EMAIL\")\n",
    "\n",
    "        def send_email(to_email: str, subject: str, body: str):\n",
    "            print(f\"[EMAIL] Attempting to send email to: {to_email}\")\n",
    "            print(f\"[EMAIL] Subject: {subject}\")\n",
    "            if not mailgun_api_key or not mailgun_domain:\n",
    "                error_msg = \"Email not configured - missing MAILGUN_API_KEY or MAILGUN_DOMAIN in environment\"\n",
    "                print(f\"[EMAIL ERROR] {error_msg}\")\n",
    "                return error_msg\n",
    "            print(\"[EMAIL] Sending via Mailgun API...\")\n",
    "            try:\n",
    "                url = f\"https://api.mailgun.net/v3/{mailgun_domain}/messages\"\n",
    "                response = requests.post(\n",
    "                    url,\n",
    "                    auth=(\"api\", mailgun_api_key),\n",
    "                    data={\"from\": mailgun_from, \"to\": to_email, \"subject\": subject, \"text\": body}\n",
    "                )\n",
    "                print(f\"[EMAIL] Response status: {response.status_code}\")\n",
    "                if response.status_code == 200:\n",
    "                    success_msg = \"Email sent successfully\"\n",
    "                    print(f\"[EMAIL SUCCESS] {success_msg}\")\n",
    "                    return success_msg\n",
    "                error_msg = f\"Email failed: {response.status_code} - {response.text}\"\n",
    "                print(f\"[EMAIL ERROR] {error_msg}\")\n",
    "                return error_msg\n",
    "            except Exception as e:\n",
    "                error_msg = f\"Email exception: {str(e)}\"\n",
    "                print(f\"[EMAIL EXCEPTION] {error_msg}\")\n",
    "                return error_msg\n",
    "\n",
    "        email_tool = Tool(\n",
    "            name=\"send_email\",\n",
    "            func=send_email,\n",
    "            description=\"Use this tool to send an email. Input should be: to_email, subject, body separated by ||| like: user@example.com|||Subject|||Body text\"\n",
    "        )\n",
    "\n",
    "        return browser_tools + file_tools + [push_tool, search_tool, python_repl, wiki_tool, email_tool]\n",
    "\n",
    "    def planner(self, state: State) -> Dict[str, Any]:\n",
    "        if state.get(\"planning_complete\"):\n",
    "            return {\"planning_complete\": True}\n",
    "\n",
    "        questions_asked = state.get(\"clarifying_questions_asked\", 0)\n",
    "        max_questions = 3\n",
    "\n",
    "        system_message = f\"\"\"You are a planning agent that clarifies user requests before work begins.\n",
    "Your job is to ensure the user's intention and success criteria are crystal clear.\n",
    "\n",
    "If anything is ambiguous or unclear about the request or success criteria, ask ONE clarifying question in a conversational, friendly tone.\n",
    "Once you're satisfied everything is clear, set ready_to_proceed to True.\n",
    "\n",
    "You have asked {questions_asked} out of a maximum of {max_questions} clarifying questions.\n",
    "If you've reached the maximum, you must proceed even if things aren't perfectly clear.\n",
    "\n",
    "Guidelines:\n",
    "- Ask questions conversationally, not in an itemized list\n",
    "- Focus on understanding the user's true intent\n",
    "- Consider what information would help the worker succeed\n",
    "- If the request is already clear, proceed immediately\n",
    "\n",
    "Success criteria provided by user: {state[\"success_criteria\"]}\n",
    "\"\"\"\n",
    "\n",
    "        messages = state[\"messages\"]\n",
    "        found_system_message = False\n",
    "        for message in messages:\n",
    "            if isinstance(message, SystemMessage):\n",
    "                message.content = system_message\n",
    "                found_system_message = True\n",
    "\n",
    "        if not found_system_message:\n",
    "            messages = [SystemMessage(content=system_message)] + messages\n",
    "\n",
    "        result = self.planner_llm_with_output.invoke(messages)\n",
    "\n",
    "        if result.ready_to_proceed or questions_asked >= max_questions:\n",
    "            return {\"planning_complete\": True}\n",
    "\n",
    "        if result.clarification_question:\n",
    "            return {\n",
    "                \"messages\": [AIMessage(content=result.clarification_question)],\n",
    "                \"clarifying_questions_asked\": questions_asked + 1,\n",
    "            }\n",
    "\n",
    "        return {\"planning_complete\": True}\n",
    "\n",
    "    def planner_router(self, state: State) -> str:\n",
    "        if state.get(\"planning_complete\"):\n",
    "            return \"worker\"\n",
    "        return \"END\"\n",
    "\n",
    "    def worker(self, state: State) -> Dict[str, Any]:\n",
    "        system_message = f\"\"\"You are a helpful assistant that can use tools to complete tasks.\n",
    "You keep working on a task until either you have a question or clarification for the user, or the success criteria is met.\n",
    "You have many tools to help you, including tools to browse the internet, navigating and retrieving web pages.\n",
    "You have a tool to run python code, but note that you would need to include a print() statement if you wanted to receive output.\n",
    "You can send emails using the send_email tool.\n",
    "The current date and time is {datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")}\n",
    "\n",
    "This is the success criteria:\n",
    "{state[\"success_criteria\"]}\n",
    "You should reply either with a question for the user about this assignment, or with your final response.\n",
    "If you have a question for the user, you need to reply by clearly stating your question. An example might be:\n",
    "\n",
    "Question: please clarify whether you want a summary or a detailed answer\n",
    "\n",
    "If you've finished, reply with the final answer, and don't ask a question; simply reply with the answer.\n",
    "\"\"\"\n",
    "\n",
    "        if state.get(\"feedback_on_work\"):\n",
    "            system_message += f\"\"\"\n",
    "Previously you thought you completed the assignment, but your reply was rejected because the success criteria was not met.\n",
    "Here is the feedback on why this was rejected:\n",
    "{state[\"feedback_on_work\"]}\n",
    "With this feedback, please continue the assignment, ensuring that you meet the success criteria or have a question for the user.\"\"\"\n",
    "\n",
    "        found_system_message = False\n",
    "        messages = state[\"messages\"]\n",
    "        for message in messages:\n",
    "            if isinstance(message, SystemMessage):\n",
    "                message.content = system_message\n",
    "                found_system_message = True\n",
    "\n",
    "        if not found_system_message:\n",
    "            messages = [SystemMessage(content=system_message)] + messages\n",
    "\n",
    "        response = self.worker_llm_with_tools.invoke(messages)\n",
    "\n",
    "        return {\"messages\": [response]}\n",
    "\n",
    "    def worker_router(self, state: State) -> str:\n",
    "        last_message = state[\"messages\"][-1]\n",
    "\n",
    "        if hasattr(last_message, \"tool_calls\") and last_message.tool_calls:\n",
    "            return \"tools\"\n",
    "        return \"evaluator\"\n",
    "\n",
    "    def format_conversation(self, messages: List[Any]) -> str:\n",
    "        conversation = \"Conversation history:\\n\\n\"\n",
    "        for message in messages:\n",
    "            if isinstance(message, HumanMessage):\n",
    "                conversation += f\"User: {message.content}\\n\"\n",
    "            elif isinstance(message, AIMessage):\n",
    "                text = message.content or \"[Tools use]\"\n",
    "                conversation += f\"Assistant: {text}\\n\"\n",
    "        return conversation\n",
    "\n",
    "    def evaluator(self, state: State) -> State:\n",
    "        last_response = state[\"messages\"][-1].content\n",
    "\n",
    "        system_message = \"\"\"You are an evaluator that determines if a task has been completed successfully by an Assistant.\n",
    "Be practical and generous in your evaluation. If the Assistant has made a genuine attempt and completed the main objective, accept it.\n",
    "Only reject if critical parts are clearly missing or the Assistant is asking questions.\"\"\"\n",
    "\n",
    "        user_message = f\"\"\"You are evaluating a conversation between the User and Assistant.\n",
    "\n",
    "The entire conversation history:\n",
    "{self.format_conversation(state[\"messages\"])}\n",
    "\n",
    "Success criteria: {state[\"success_criteria\"]}\n",
    "\n",
    "Final response from Assistant: {last_response}\n",
    "\n",
    "Evaluation guidelines:\n",
    "- If files were created using tools, assume they were successful\n",
    "- If notifications were sent using tools, assume they were successful  \n",
    "- If the Assistant completed the main task, accept it even if minor details are imperfect\n",
    "- Only mark user_input_needed=True if the Assistant explicitly asks a question\n",
    "- Give the Assistant the benefit of the doubt on tool usage\n",
    "- Accept reasonable efforts that meet the core success criteria\n",
    "\"\"\"\n",
    "        if state[\"feedback_on_work\"]:\n",
    "            user_message += f\"\\nPrevious feedback: {state['feedback_on_work']}\\nIf the Assistant is stuck in a loop or repeating, mark user_input_needed=True.\"\n",
    "\n",
    "        evaluator_messages = [\n",
    "            SystemMessage(content=system_message),\n",
    "            HumanMessage(content=user_message),\n",
    "        ]\n",
    "\n",
    "        eval_result = self.evaluator_llm_with_output.invoke(evaluator_messages)\n",
    "        new_state = {\n",
    "            \"messages\": [{\"role\": \"assistant\", \"content\": f\"Evaluator Feedback: {eval_result.feedback}\"}],\n",
    "            \"feedback_on_work\": eval_result.feedback,\n",
    "            \"success_criteria_met\": eval_result.success_criteria_met,\n",
    "            \"user_input_needed\": eval_result.user_input_needed,\n",
    "        }\n",
    "        return new_state\n",
    "\n",
    "    def route_based_on_evaluation(self, state: State) -> str:\n",
    "        if state[\"success_criteria_met\"] or state[\"user_input_needed\"]:\n",
    "            return \"END\"\n",
    "        return \"worker\"\n",
    "\n",
    "    async def build_graph(self):\n",
    "        graph_builder = StateGraph(State)\n",
    "\n",
    "        graph_builder.add_node(\"planner\", self.planner)\n",
    "        graph_builder.add_node(\"worker\", self.worker)\n",
    "        graph_builder.add_node(\"tools\", ToolNode(tools=self.tools))\n",
    "        graph_builder.add_node(\"evaluator\", self.evaluator)\n",
    "\n",
    "        graph_builder.add_conditional_edges(\"planner\", self.planner_router, {\"worker\": \"worker\", \"END\": END})\n",
    "        graph_builder.add_conditional_edges(\"worker\", self.worker_router, {\"tools\": \"tools\", \"evaluator\": \"evaluator\"})\n",
    "        graph_builder.add_edge(\"tools\", \"worker\")\n",
    "        graph_builder.add_conditional_edges(\"evaluator\", self.route_based_on_evaluation, {\"worker\": \"worker\", \"END\": END})\n",
    "        graph_builder.add_edge(START, \"planner\")\n",
    "\n",
    "        self.graph = graph_builder.compile(checkpointer=self.memory)\n",
    "\n",
    "    async def run_superstep(self, message, success_criteria, history):\n",
    "        config = {\n",
    "            \"configurable\": {\"thread_id\": self.sidekick_id},\n",
    "            \"recursion_limit\": 50\n",
    "        }\n",
    "\n",
    "        state = {\n",
    "            \"messages\": message,\n",
    "            \"success_criteria\": success_criteria or \"The answer should be clear and accurate\",\n",
    "            \"feedback_on_work\": None,\n",
    "            \"success_criteria_met\": False,\n",
    "            \"user_input_needed\": False,\n",
    "            \"clarifying_questions_asked\": 0,\n",
    "            \"planning_complete\": False,\n",
    "        }\n",
    "        result = await self.graph.ainvoke(state, config=config)\n",
    "        user = {\"role\": \"user\", \"content\": message}\n",
    "        reply = {\"role\": \"assistant\", \"content\": result[\"messages\"][-2].content}\n",
    "        feedback = {\"role\": \"assistant\", \"content\": result[\"messages\"][-1].content}\n",
    "        return history + [user, reply, feedback]\n",
    "\n",
    "    def cleanup(self):\n",
    "        if self.browser:\n",
    "            try:\n",
    "                loop = asyncio.get_running_loop()\n",
    "                loop.create_task(self.browser.close())\n",
    "                if self.playwright:\n",
    "                    loop.create_task(self.playwright.stop())\n",
    "            except RuntimeError:\n",
    "                asyncio.run(self.browser.close())\n",
    "                if self.playwright:\n",
    "                    asyncio.run(self.playwright.stop())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35f1a18d",
   "metadata": {},
   "source": [
    "### Initialize Sidekick and Gradio Interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a13ebcfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sidekick = None\n",
    "\n",
    "\n",
    "async def init_sidekick():\n",
    "    global sidekick\n",
    "    sidekick = Sidekick()\n",
    "    await sidekick.setup()\n",
    "\n",
    "\n",
    "def make_thread_id():\n",
    "    return str(uuid.uuid4())\n",
    "\n",
    "\n",
    "async def process_message(message, success_criteria, history, thread):\n",
    "    global sidekick\n",
    "    if not sidekick:\n",
    "        await init_sidekick()\n",
    "    return await sidekick.run_superstep(message, success_criteria, history)\n",
    "\n",
    "\n",
    "async def reset():\n",
    "    return \"\", \"\", None, make_thread_id()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06b08067",
   "metadata": {},
   "source": [
    "### Launch Sidekick UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f1197f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks(theme=gr.themes.Default(primary_hue=\"emerald\")) as demo:\n",
    "    gr.Markdown(\"## Sidekick Personal Co-worker\")\n",
    "    thread = gr.State(make_thread_id())\n",
    "\n",
    "    with gr.Row():\n",
    "        chatbot = gr.Chatbot(label=\"Sidekick\", height=300, type=\"messages\")\n",
    "    with gr.Group():\n",
    "        with gr.Row():\n",
    "            message = gr.Textbox(show_label=False, placeholder=\"Your request to your sidekick\")\n",
    "        with gr.Row():\n",
    "            success_criteria = gr.Textbox(show_label=False, placeholder=\"What are your success criteria?\")\n",
    "    with gr.Row():\n",
    "        reset_button = gr.Button(\"Reset\", variant=\"stop\")\n",
    "        go_button = gr.Button(\"Go!\", variant=\"primary\")\n",
    "    message.submit(process_message, [message, success_criteria, chatbot, thread], [chatbot])\n",
    "    success_criteria.submit(process_message, [message, success_criteria, chatbot, thread], [chatbot])\n",
    "    go_button.click(process_message, [message, success_criteria, chatbot, thread], [chatbot])\n",
    "    reset_button.click(reset, [], [message, success_criteria, chatbot, thread])\n",
    "\n",
    "\n",
    "demo.launch(debug=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
